#!/usr/local/bin/python3.7
import pyaudio
import wave
import os
import sys
import threading
from audiosplit import splitwav
from asr import asrwavs
from multiprocessing import Pipe
import tkinter as tk
from tkinter.filedialog import asksaveasfile

CHUNK = 512
FORMAT = pyaudio.paInt16
CHANNELS = 1
RATE = 16*1024

# 默认180秒保存一次数据，并开始识别
RECORD_SECONDS = 180

# record_stat :0-continue,1-save and continue,2-save and end
record_stat = 0

split_tmp_ext = 1

output_texts = []

class myAsr (threading.Thread):
	def __init__(self,tkWindow):
		threading.Thread.__init__(self)
		self.tkWindow = tkWindow
	def run(self):
		global split_tmp_ext
		global precv
		global output_texts
		while True:
			try:
				filename = precv.recv()
			except:
				print('end.')
				break
			print(f'开始识别：{filename}')
			names = splitwav(filename.strip(), split_tmp_ext)
			split_tmp_ext += 1
			if len(names)==0:
				print('wave file split fail.')
				os.exit(0)
			res = asrwavs(names)
			print( res )
			output_texts.append(res)
			#self.tkWindow.text1.insert(tk.END,res)
			#清理临时文件和文件夹
			for i in range(len(names)):
				os.unlink(names[i])
			os.rmdir( os.path.dirname(names[0]) )

class myThread (threading.Thread):
	def __init__(self, tkWindow):
		threading.Thread.__init__(self)
		self.tkWindow = tkWindow
	def run(self):
		threads = recordaudio()
		for i in range(len(threads)):
			threads[i].join()

def control_by_key():
	global record_stat
	while True:
		print('按回车键保存当前音频;')
		print('输入\'q\'再按回车键保存并结束程序。')
		s = sys.stdin.readline()
		if s.strip() == '':
			record_stat = 1
			print('你按了回车键')
		elif s.strip() == 'q':
			record_stat = 2
			print('你输入了\'q\'，退出程序')
			break
			
def savewav( frames:list, filename:str ):
	wf = wave.open(filename, 'wb')
	wf.setnchannels(CHANNELS)
	wf.setsampwidth(pyaudio.get_sample_size(FORMAT))
	wf.setframerate(RATE)
	wf.writeframes(b''.join(frames))
	wf.close()

def recordaudio() -> list:
	global record_stat
	global port
	
	p = pyaudio.PyAudio()
	stream = p.open(format=FORMAT,
				channels=CHANNELS,
				rate=RATE,
				input=True,
				frames_per_buffer=CHUNK)

	print("recording...")

	frames = []
	idx = 1
	threads = []
	global psend
	while record_stat < 2:
		for i in range(0, int(RATE / CHUNK * RECORD_SECONDS)):
			if record_stat > 0:
				break
			data = stream.read(CHUNK)
			frames.append(data)
		if record_stat == 1:
			record_stat = 0
		savewav(frames, f'rec/a{idx:04}.wav')
		print(f'save wav rec/a{idx:04}.wav')
		psend.send(f'rec/a{idx:04}.wav\n')
		idx+=1
		frames.clear()
	psend.close()
	
	# Stop and close the stream 
	stream.stop_stream()
	stream.close()
	# Terminate the PortAudio interface
	p.terminate()
	return threads
	
class myWindow:
	def __init__(self, title='实时语音识别'):
		self.root = tk.Tk(className=title)

	def packBtn(self):
		self.f1=tk.Frame(self.root)
		self.f1.pack(side='top')
		self.f2=tk.Frame(self.root)
		self.f2.pack(side='bottom')
		self.btnStart=tk.Button(self.f1, command=self.start_record, width=6, height=3)
		self.btnStart.pack(padx=10, side='left')
		self.btnStart['text'] = '开始'

		self.btnDuanju=tk.Button(self.f1, command=self.duanju, width=6, height=3)
		self.btnDuanju.pack(padx=10, side='left')
		self.btnDuanju['text'] = '断句'
		self.btnDuanju.config(state=tk.DISABLED)
		
		self.btnEnd = tk.Button(self.f1, command=self.end_record, width=6, height=3)
		self.btnEnd.pack(padx=10, side='left')
		self.btnEnd['text'] = '结束'
		self.btnEnd.config(state=tk.DISABLED)
		
		self.btnSave = tk.Button(self.f1, command=self.save_result, width=6, height=3)
		self.btnSave.pack(padx=10, side='left')
		self.btnSave['text'] = '保存'
		self.btnSave.config(state=tk.DISABLED)
		
		self.text1 = tk.Text(self.f2)
		self.scroll1 = tk.Scrollbar( self.f2, width=20 )
		
		self.scroll1.pack(side='right', fill=tk.Y)
		self.text1.pack(  fill=tk.BOTH,expand=True,side='left' )
		
		self.scroll1['command'] = self.text1.yview
		self.text1['yscrollcommand'] = self.scroll1.set
		
	def loop(self):
		self.root.resizable(False, False)   #禁止修改窗口大小
		self.packBtn()
		self.root.mainloop()

	def start_record(self):
		global psend
		global precv
		psend, precv = Pipe()
		self.thread_rec = myThread(self)
		self.thread_rec.start()
		self.thread_asr = myAsr(self)
		self.thread_asr.start()
		self.btnStart.config(state=tk.DISABLED)
		self.btnDuanju.config(state=tk.NORMAL)
		self.btnEnd.config(state=tk.NORMAL)
		
	def duanju(self):
		global record_stat
		record_stat = 1
		
	def end_record(self):
		global record_stat
		record_stat = 2
		self.thread_rec.join()
		self.thread_asr.join()
		self.text1.insert( tk.END, ''.join(output_texts) )
		self.btnStart.config(state=tk.NORMAL)
		self.btnDuanju.config(state=tk.DISABLED)
		self.btnEnd.config(state=tk.DISABLED)
		self.btnSave.config(state=tk.NORMAL)
		
	def save_result(self):
		f = asksaveasfile()
		for i in range(len(output_texts)):
			f.write(output_texts[i])
		f.close()
		
if __name__=='__main__':
	try:
		os.mkdir('rec')
	except:
		None
	w = myWindow()
	w.loop()

	# output_name = 'output.txt'
	# if len(sys.argv)==2:
		# output_name = sys.argv[1]
	# fp = open( output_name, 'w' )
	# for i in range(len(output_texts)):
		# fp.write(output_texts[i])
	# fp.close()
